import os
import numpy as np
import torch
import tqdm
from torch.utils.data import DataLoader

def get_data(data_path='data/tang.npz'):
    if os.path.exists(data_path):
        datas = np.load(data_path, allow_pickle=True)      # 加载数据
        data = datas['data']  # numpy.ndarray
        word2ix = datas['word2ix'].item()   # dic
        ix2word = datas['ix2word'].item()  # dic
        return data, word2ix, ix2word

def make_dataloader(data_path='data/tang.npz', batch_size=10):
    data, word2ix, ix2word = get_data(data_path = data_path)
    data = torch.from_numpy(data)
    dataloader = DataLoader(data, batch_size=batch_size, shuffle=True)
    return dataloader

def test():
    _, word2ix, ix2word = get_data(data_path='data/tang.npz')
    dataloader = make_dataloader(batch_size=2)
    for i, data in tqdm.tqdm(enumerate(dataloader)): 
        data = data.long().contiguous()
        print(data.shape)
        input, target = data[:-1, :], data[1:, :]
        input_word = []
        target_word = []
        #print(input.shape)
        #print(target.shape)
        for i in range(input.shape[1]):
            input_word.append(ix2word[int(input[0,i])])
        #print(input_word)
        for i in range(target.shape[1]):
            target_word.append(ix2word[int(target[0,i])])
        #print(target_word)
        break

# test()